import streamlit as st
import pandas as pd
import logging
import numpy as np
import time
import io
import os
import subprocess
from project.invoke_functions import convert_excel_to_json, start_scan
from project.colored_logging import console_table, rich_startup_info, art_startup_Info
import project.script_launcher as script_launcher
import threading
import contextlib


# 创建一个自定义的日志处理程序
class StreamlitLogHandler(logging.Handler):
    def __init__(self):
        super().__init__()
        self.log_content = ""

    def emit(self, record):
        # 使用logging模块的默认格式化器获取时间戳
        formatter = logging.Formatter()
        log_message = f'{formatter.formatTime(record)} - < {record.levelname} >: {record.msg}'
        self.log_content += log_message + "\n"


# 配置日志记录器，使用自定义的日志处理程序
logging.basicConfig(level=logging.DEBUG)
log_handler = StreamlitLogHandler()
logging.getLogger().addHandler(log_handler)

MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB

def read_excel_to_dict(file):
    """
    读取上传的Excel文件的第一个工作表，并转换为字典。

    参数：
    file：上传的文件对象。

    返回：
    如果读取成功，返回包含数据的字典；否则返回None。
    """
    try:
        df = pd.read_excel(file, sheet_name=0)
        data_dict = df.to_dict(orient='records')
        return data_dict
    except pd.errors.ExcelFileError as e:
        st.error(f"在数据展示中：文件读取失败: {e}。请确保文件格式正确且未损坏。")
        return None


def create_download_button(data_dict, export_format):
    """
    创建下载按钮。

    参数：
    data_dict：修改后的字典数据。
    export_format：导出的文件格式。
    """
    if export_format == 'xlsx':
        df = pd.DataFrame(data_dict)
        buffer = io.BytesIO()
        df.to_excel(buffer, index=False)
        buffer.seek(0)
        output_data = buffer.getvalue()
        mime_type = 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
        file_name = "sheet1_data.xlsx"
        label = f'下载数据'
    elif export_format == 'xls':
        df = pd.DataFrame(data_dict)
        buffer = io.BytesIO()
        df.to_excel(buffer, index=False)
        buffer.seek(0)
        output_data = buffer.getvalue()
        mime_type = 'application/vnd.ms-excel'
        file_name = "sheet1_data.xls"
        label = f'下载数据'

    st.download_button(
        label=label,
        data=output_data,
        file_name=file_name,
        mime=mime_type
    )


def start_script_and_actions(script_path, uploaded_file, line=None):
    """
    启动脚本并执行相关操作。

    参数：
    script_path：脚本路径。
    uploaded_file：上传的文件对象。
    line：要处理的行号。
    """
    with log_handler.lock:
        print()
        logging.info("脚本线程准备启动。")
        # 检查是否已有脚本线程在运行
        existing_threads = [t for t in threading.enumerate() if t.name == "script_thread"]
        if len(existing_threads) > 0:
            logging.info("已有脚本线程在运行，请等待其完成或手动终止后再启动新的脚本线程。")
        # 如果没有脚本线程在运行，则启动新线程
        if not existing_threads:
            thread = threading.Thread(target=script_launcher.start_script, args=(script_path,), name="script_thread")
            thread.start()
            logging.info("脚本线程成功启动。")
            # 等待一小段时间确保脚本启动
            time.sleep(2)
        # 无论是否有脚本线程，都执行后续操作
        convert_excel_to_json(uploaded_file)
        start_scan()



def capture_console_output(func):
    """
    装饰器函数，用于捕获被装饰函数执行过程中的控制台输出内容，
    并将其添加到日志处理程序的log_content中以便后续展示在tab2中。
    """
    def wrapper(*args, **kwargs):
        console_output = io.StringIO()
        with contextlib.redirect_stdout(console_output):
            result = func(*args, **kwargs)
            console_text = console_output.getvalue()
            if console_text:
                with log_handler.lock:
                    record = logging.makeLogRecord({
                        'name': 'captured_output',
                        'levelno': logging.INFO,
                        'levelname': 'INFO',
                        'msg': console_text.strip(),
                    })
                    log_handler.emit(record)
        return result
    return wrapper


def capture_script_logs(script_path):
    try:
        process = subprocess.Popen(
            f'python {script_path}',
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT
        )
        for line in process.stdout:
            with log_handler.lock:
                try:
                    # 尝试对日志行进行解码，使用utf-8编码，如果出现错误则尝试其他编码
                    decoded_line = line.decode('utf-8')
                except UnicodeDecodeError:
                    decoded_line = line.decode('latin-1')
                record = logging.makeLogRecord({
                    'name': 'script_output',
                    'levelno': logging.INFO,
                    'levelname': 'INFO',
                    'msg': decoded_line.strip(),
                })
                log_handler.emit(record)
        process.wait()
    except Exception as e:
        logging.error(f"捕获脚本日志时出错：{e}")


def update_log_display():
    # 更新日志展示
    log_content_to_show = []
    if not log_handler.log_content:
        log_content_to_show.append("无日志数据信息")
    else:
        lines = log_handler.log_content.split('\n')
        formatted_lines = []
        for line in lines:
            parts = line.split(' - < ')
            if len(parts) == 2:
                time_part, rest = parts
                level_part, msg_part = rest.split(' >: ')
                formatted_lines.append(f'{time_part} < {level_part} >: {msg_part}')
            else:
                formatted_lines.append(line)
        log_content_to_show.extend(formatted_lines)
    log_content_text = "<br>".join(log_content_to_show)
    st.markdown(f'<div class="log-info-content"><p style="width:100%; overflow-wrap:break-word; font-size: 13px;">{log_content_text}</p></div>',
                unsafe_allow_html=True)
    

def convert_all_to_str(data):
    """
    数据格式化
    参数：
    data：输入的数据，可以是 DataFrame、列表或字典。

    返回：
    """
    if isinstance(data, dict):
        result = {}
        for key, value in data.items():
            result[key] = convert_all_to_str(value)
        return result
    elif isinstance(data, list):
        return [convert_all_to_str(item) for item in data]
    elif isinstance(data, float) and np.isnan(data):
        return data
    else:
        if isinstance(data, float):
            data = int(data)
        return str(data)

def add_data_type(data):
    """
    此函数接收不同类型的数据（如 DataFrame、列表、字典），为其元素添加数据类型，并重新存储。
    对于 DataFrame，会将每列元素的数据类型添加到元素中；对于列表和字典，会为元素添加数据类型。

    参数：
    data：输入的数据，可以是 DataFrame、列表或字典。

    返回：
    存储了添加数据类型元素的新数据结构，与输入数据类型相同。
    """
    if isinstance(data, pd.DataFrame):
        result_df = data.copy()
        for column in result_df.columns:
            result_df[column] = result_df[column].apply(lambda x: str(x) + ' (' + str(type(x).__name__) + ')')
        console_table(result_df)
    elif isinstance(data, list):
        result_list = []
        for item in data:
            if isinstance(item, dict):
                result_list.append(add_data_type(item))
            else:
                result_list.append(str(item) + ' (' + str(type(item).__name__) + ')')
        console_table(result_list)
    elif isinstance(data, dict):
        result_dict = {}
        for key, value in data.items():
            result_dict[key] = str(value) + ' (' + str(type(value).__name__) + ')'
        console_table(result_dict)
    else:
        raise ValueError("输入的数据类型不支持，仅支持 DataFrame、列表和字典。")
    

def find_matching_key_value(record, target_key):
    """
    在给定的字典中查找与目标键匹配的值。
    先查找 ip 相关键，再查找 port 相关键，将找到的 ip 和 port 组合成元组，与 target_key 进行比较。

    参数:
    record: 要查找的字典。
    target_key: 目标键的值，为元组 (ip, port)。

    返回:
    如果找到匹配的值，返回该值；否则返回 None。
    """
    ip_keys = ['ip', 'IP', 'iP', 'Ip', '主机']
    port_keys = ['port', 'Port', 'PORT', '端口']
    possible_ip = None
    possible_port = None
    # 查找 ip 相关键对应的可能值
    for ip_key in ip_keys:
        if ip_key in record:
            possible_ip = record[ip_key]
            break
    # 查找 port 相关键对应的可能值
    for port_key in port_keys:
        if port_key in record:
            possible_port = record[port_key]
            break
    # 组合可能的 ip 和 port 值
    possible_pair = (possible_ip, str(possible_port))
    # 比较组合值与目标键
    if possible_pair == target_key:
        return possible_pair    
    else:
        return None

def append_non_matching_data(index, target_dict, source_dict, keys_to_exclude):
    """
    将源字典中除了指定键之外的其他键值对追加到目标字典中。

    参数：
    target_dict: 目标字典，数据将追加到此字典中。
    source_dict: 源字典，从此字典获取要追加的数据。
    keys_to_exclude: 要排除的键列表。
    """
    for key in source_dict:
        if key not in keys_to_exclude:
            target_dict[key] = source_dict[key]
    return target_dict


def data_consolidation(data_int_path, data_dict_path):
    """
    基于已知条件，检索 data_int_path (对应 Excel 文件转换后的字典数据) 和 data_dict_path 两个字典中特定键对应相同值的记录，
    按照映射关系将 data_int_path 字典中除特定键之外的其他列数据追加到 data_dict_path 字典中。

    参数：
    data_int_path: 扫描任务完成导出的 excel 文件路径。
    data_dict_path: 初始 upload 上传的文件。

    返回：
    整合后的字典。
    """
    try:
        if not os.path.exists(data_int_path):
            st.error(f"指定的整合数据文件路径 {data_int_path} 不存在，请检查路径是否正确。")
            return None
        data_int = pd.read_excel(data_int_path)
        data_records = data_int.to_dict(orient='records')
        logging.info("正在获取扫描结果")

        consolidated_data = data_dict_path.copy()
        consolidated_str = convert_all_to_str(consolidated_data)
        data_records_str = convert_all_to_str(data_records)

        logging.info("数据格式化完成")
        logging.info("数据正在匹配中")
        for item_data_int in data_records_str:
            fixed_ip = item_data_int.get('Host', None)
            fixed_port = item_data_int.get('Port', None)
            target_pair = (fixed_ip, str(fixed_port))
            index = 0
            for item_data_dict in consolidated_str:
                match = find_matching_key_value(item_data_dict, target_pair)
                try:
                    if match:
                        append_non_matching_data(index,item_data_dict, item_data_int, ['Host', 'Port'])
                        index += 1
                except Exception as e:
                    st.error(f"数据整合过程出错: {e}")
                    return None
        logging.info("成功获取扫描结果")
        return consolidated_str
    except pd.errors.ExcelFileError as e:
        st.error(f"整合数据文件读取失败: {e}。请确保文件格式正确且未损坏，或者尝试使用不同的Excel文件格式。")
        return None
    except Exception as e:
        st.error(f"数据整合过程出错: {e}")
        return None
    
def get_dataframe_to_show(upload_dict, consolidated_data_list, start_scan=False):
    if start_scan:
        if consolidated_data_list:
            upload_dict = None
            if isinstance(consolidated_data_list, list) and len(consolidated_data_list) > 0:
                logging.info("Returnint consolidated List:")
                console_table(consolidated_data_list)
                return consolidated_data_list
            else:
                print("consolidated_data_dict is empty or None.")
                return None
        else:
            logging.info("Returning data_dict List:")
            console_table(upload_dict)
            return upload_dict
    elif upload_dict:
        return upload_dict
    logging.error("No valid data to show.")
    return None

def main():
    st.set_page_config(layout="wide")
    st.title("Scan Network Excel")

    # 文件上传在选项卡外面
    uploaded_file = st.file_uploader(label="Please upload an Excel file.", type=["xlsx", "xls"], key="file_uploader")

    script_thread = None

    # 创建两个选项卡
    tab1, tab2 = st.tabs(["数据展示", "日志信息"])

    # 数据展示选项卡
    with tab1:
        upload_dict = None
        consolidated_data_list = {}
        # Excel文件处理部分
        if uploaded_file is not None:
            if uploaded_file.size > MAX_FILE_SIZE:
                st.error("在数据展示中：文件过大，请上传小于10MB的文件。")
                return

            file_name = uploaded_file.name
            if file_name.endswith('.xlsx'):
                export_format = 'xlsx'
            elif file_name.endswith('.xls'):
                export_format = 'xls'
            else:
                st.error("在数据展示中：上传的文件不是有效的Excel文件，请重新上传。")
                export_format = None

            if export_format:
                with st.spinner('正在读取文件...'):
                    upload_dict = read_excel_to_dict(uploaded_file)
                if upload_dict is not None:
                    current_dir = os.path.dirname(os.path.abspath(__file__))
                    data_int_path = os.path.join(current_dir, 'data', 'nmap_results.xlsx')
                    start_scan = False
                    if st.button("开始扫描"):
                        start_scan = True
                        art_startup_Info("STARTUP")
                        current_dir = os.path.dirname(os.path.abspath(__file__))
                        script_path = os.path.join(current_dir, 'project', 'nmap_scan.py')
                        script_thread = start_script_and_actions(script_path, uploaded_file)
                        
                        consolidated_data_list = data_consolidation(data_int_path, upload_dict)
                    list_to_show = get_dataframe_to_show(upload_dict, consolidated_data_list, start_scan)
                    if list_to_show is not None:
                        df_to_show = pd.DataFrame(list_to_show)
                        # 使用st.data_editor创建一个可编辑的数据表格
                        edited_df = st.data_editor(df_to_show, num_rows="dynamic", key="data_editor")
                        # 将编辑后的数据转换回字典格式
                        # upload_dict = edited_df.to_dict(orient='records')

                    # 导出按钮，使用更新后的数据字典创建下载按钮
                    create_download_button(edited_df, export_format)
                else:
                    st.warning("在数据展示中：无法读取上传的文件，请检查文件是否损坏或是否为支持的Excel格式。")
            else:
                st.error("在数据展示中：上传的文件不是有效的Excel文件，请重新上传。")
        else:
            st.info("在数据展示中：请上传一个Excel文件。")

    with tab2:
        # 使用st.markdown创建一个隐藏的div来注入样式，确保样式不会被直接显示在页面上
        st.markdown('<div style="display:none;"><style>.log-info-content {height: 500px; background-color: black; color: white; width: 100%; overflow-y: auto; padding: 10px;}</style></div>',
                    unsafe_allow_html=True)
        # 初始日志展示
        update_log_display()

    # 通过css控制表格宽度
    st.markdown('<style>.dataframe{width:100%!important;}</style>', unsafe_allow_html=True)

    # 在函数结束前检查脚本线程是否已结束，如果没有结束，则等待其结束
    if script_thread and script_thread.is_alive():
        script_thread.join()


if __name__ == "__main__":
    main()
